# -*- coding: utf-8 -*-
"""
@Author ：mengying
@Date   ：2024/3/20 9:49
@Email  : 652044581@qq.com
@Desc   : mask和图片的抠图
"""

import numpy as np
from PIL import Image
import imutils


class Mask2Image:
    """抠图转换成png 含Alpha通道"""

    def __init__(self, mask: np.ndarray, image: np.ndarray):
        self.mask = mask
        self.image = image
        self.preprocess()

    @classmethod
    def read_filePath(cls, filePath):
        return cv2.imread(filePath)

    def preprocess(self):
        """图片进行预处理, 腐蚀5个像素， 避免分割线突兀"""
        kernel = np.ones((5, 5), np.uint8)
        self.mask = cv2.erode(self.mask, kernel, iterations=1)

    def transform_png(self) -> np.ndarray:
        self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2BGRA)
        return cv2.bitwise_and(self.image, self.image, mask=self.mask)


class ReplaceBackground:
    """替换图片的背景"""

    def __init__(self, background_color=None, size=None, scale=None):
        self.background_color = background_color or (255, 255, 255)  # 背景默认白色
        self.size = size or (800, 800)  # 默认800X800
        self.scale = scale or 1  # 占比默认100%

    @classmethod
    def cv2PIL(cls, image):
        """opencv和 pil的格式转换"""
        return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA))

    @classmethod
    def PIL2cv(cls, image):
        return cv2.cvtColor(np.asarray(image), cv2.COLOR_RGBA2BGR)

    def get_background(self) -> np.ndarray:
        """生成颜色大小的背景图片"""
        background = np.ones((self.size[0], self.size[1], 3), dtype=np.uint8)
        background[:, :, 0] = self.background_color[0]  # 蓝色通道B
        background[:, :, 1] = self.background_color[1]  # 绿色通道G
        background[:, :, 2] = self.background_color[2]  # 红色通道R
        return background

    def get_bounding_rect(self, image: np.ndarray) -> np.ndarray:
        """找到最大轮廓的最小矩形"""
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 1, 225)
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # 找到轮廓中最大面积的然后返回
        max_area = 0
        index = None
        for i, cnt in enumerate(contours):
            x, y, w, h = cv2.boundingRect(cnt)
            area = w * h
            if area > max_area:
                max_area = area
                index = i

        if index is not None:
            x, y, w, h = cv2.boundingRect(contours[index])
            return image[y:y + h, x:x + w, :]
        return image

    @classmethod
    def replace_background_color(cls, image: np.ndarray, color) -> Image.Image:
        """替换背景图片， BGRA格式"""
        image_bg = Image.new('RGB', (image.shape[1], image.shape[0]), (color[2], color[1], color[0]))
        image_with_alpha = cls.cv2PIL(image)
        alpha = image_with_alpha.split()[-1]
        mask = Image.eval(alpha, lambda x: 255 if x > 0 else 0)
        image_bg.paste(image_with_alpha, mask=mask)
        return image_bg

    def transform(self, image: np.ndarray) -> Image.Image:
        background = self.get_background()
        image_bg = self.replace_background_color(image, self.background_color)

        # 找到图像的最小外接矩形
        image = self.get_bounding_rect(self.PIL2cv(image_bg))

        # 放缩比例
        if image.shape[0] / background.shape[0] > image.shape[1] / background.shape[1]:
            image_scale = background.shape[1] * (image.shape[1] / image.shape[0])
        else:
            image_scale = background.shape[1]

        image_resize = imutils.resize(image, int(image_scale * self.scale))
        image_w = image_resize.shape[1]
        image_h = image_resize.shape[0]
        margin_w = int((background.shape[1] - image_w) / 2)
        margin_h = int((background.shape[0] - image_h) / 2)

        image_pil = self.cv2PIL(image_resize)
        background_pil = self.cv2PIL(background)
        background_pil.paste(image_pil, box=(margin_w, margin_h, margin_w + image_w, margin_h + image_h))
        return background_pil

    @classmethod
    def save_filename(cls, image, filename) -> str:
        """保存图片为路径"""
        if isinstance(image, Image.Image):
            image.save(filename)
        elif isinstance(image, np.ndarray):
            cv2.imwrite(filename, image)


if __name__ == '__main__':
    import cv2
    import numpy as np

    # 加载原始图像
    image = cv2.imread(r'C:\Users\cqax\Desktop\d8eef6ba2d4d4c29b014c4ff56b123b1.png')

    # 转换为灰度图像
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 创建一个与输入图像相同大小的全零掩膜
    mask = np.zeros(image.shape[:2], dtype=np.uint8)

    _, alpha = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)

    # 使用Canny边缘检测来检测轮廓
    edges = cv2.Canny(gray, 10, 225)

    # 寻找图像中的轮廓
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 在掩膜上绘制轮廓，将轮廓内部填充为白色
    mask = cv2.drawContours(mask, contours, -1, (255), thickness=cv2.FILLED)

    result = Mask2Image(mask, image).transform_png()

    result1 = ReplaceBackground(background_color=(255, 123, 255), size=(800, 800), scale=0.5).transform(result)

    # 设置背景为透明
    result1.show()
